"""
python fujian1_arima_gragh.py
"""

import os
import json
import pandas as pd
import matplotlib.pyplot as plt

# 输入和输出目录
input_dir = "fujian/fujian1/arima/predict_output"
output_dir = "fujian/fujian1/arima/graph"
os.makedirs(output_dir, exist_ok=True)

# 颜色定义
original_data_color = "blue"
predicted_data_color = "orange"

# 遍历每个 JSON 文件
for filename in os.listdir(input_dir):
    if filename.endswith(".json"):
        filepath = os.path.join(input_dir, filename)
        
        # 读取 JSON 文件
        with open(filepath, 'r', encoding='utf-8') as file:
            data = json.load(file)
        
        # 转换为 DataFrame
        df = pd.DataFrame(data)
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)

        # 区分原始数据和预测数据
        original_data = df.loc["2022-07-01":"2023-06-01"]
        predicted_data = df.loc["2023-07-01":"2023-09-01"]

        # 创建图形
        plt.figure(figsize=(10, 6))
        
        # 绘制原始数据
        plt.plot(original_data.index, original_data['inventory'], label="Original Data", color=original_data_color)
        
        # 绘制预测数据
        plt.plot(predicted_data.index, predicted_data['inventory'], label="Predicted Data", color=predicted_data_color, linestyle="--")
        
        # 设置图表标题和标签
        plt.title(f"Inventory Trend for {filename.split('.')[0]}")
        plt.xlabel("Date")
        plt.ylabel("Inventory")
        plt.legend()
        plt.grid()

        # 保存图表
        output_filepath = os.path.join(output_dir, f"{filename.split('.')[0]}.png")
        plt.savefig(output_filepath)
        plt.close()
        
        print(f"Graph saved: {output_filepath}")
